import os
from math import ceil
from glob import glob
from argparse import ArgumentParser

from tqdm import tqdm
import numpy as np
from affine import Affine
from osgeo import gdal
import rasterio
from rasterio.windows import Window

# tif 图像的边界
def GetExtent(infile):
    ds = gdal.Open(infile)
    geotrans = ds.GetGeoTransform()
    xsize = ds.RasterXSize
    ysize = ds.RasterYSize
    min_x,max_y = geotrans[0],geotrans[3]
    max_x,min_y = geotrans[0]+xsize*geotrans[1],geotrans[3]+ysize*geotrans[5]
    ds = None
    return min_x,max_y,max_x,min_y

def compress(path, target_path,method="LZW"): 
    """
    使用gdal进行文件压缩，
    LZW方法属于无损压缩，
    效果非常给力，4G大小的数据压缩后只有三十多M
    """
    dataset = gdal.Open(path)
    driver = gdal.GetDriverByName('GTiff')
    driver.CreateCopy(target_path, dataset, strict=1, options=["TILED=YES", "COMPRESS={0}".format(method)])
    del dataset

def RasterMosaic(file_list,outpath):
    min_x,max_y,max_x,min_y=GetExtent(file_list[0])
    for infile in tqdm(file_list):
        minx,maxy,maxx,miny = GetExtent(infile)
        min_x,min_y = min(min_x,minx),min(min_y,miny)
        max_x,max_y = max(max_x,maxx),max(max_y,maxy)
    
    in_ds = rasterio.open(file_list[0])
    
    outfile = outpath # 结果文件名，可自行修改
    geotrans = list(in_ds.transform)
    res_x, res_y = geotrans[0], -geotrans[4]
    columns = ceil((max_x-min_x)/res_x)#列数
    rows = ceil((max_y-min_y)/(res_y))#行数
    kwargs = in_ds.meta.copy() # 不能是 vrt_raster.meta.copy()
    kwargs.update({
        'height': rows,
        'width': columns,
        'transform': Affine.from_gdal(*(min_x, res_x, 0, max_y, 0, -res_y))})
    out_ds = rasterio.open(outfile, mode='w+',  **kwargs)

    for in_fn in tqdm(file_list):
        in_ds=rasterio.open(in_fn)
        height, width, trans = in_ds.height, in_ds.width, list(in_ds.transform)
        offset_y, offset_x = ceil((max_y-trans[5])/-trans[4]), ceil((trans[2]-min_x)/trans[0])
        rows = (offset_y, offset_y+height)
        cols = (offset_x, offset_x+width)
        Win = Window.from_slices(rows, cols)

        org_data = out_ds.read(window=Win)
        org_mask = out_ds.read(window=Win, masked=True).mask
        in_data = in_ds.read(masked=True)
        org_data[np.logical_not(in_data.mask)] = in_data[np.logical_not(in_data.mask)]

        # in_data = 

        # axis = (1, 2) if len(org_data.shape)>2 else (0, 1)
        # mean = np.mean(org_data, axis=axis) * 0.7 + np.mean(in_data, axis=axis) + 0.3
        # std = np.std(org_data, axis=axis) * 0.7 + np.std(in_data, axis=axis) + 0.3
        # while len(std.shape) < len(in_data.shape):
        #     np.expand_dims(std, axis=1)
        # in_data = (in_data-mean.data) / std.data
        org_data[np.logical_not(in_data.mask)] = in_data[np.logical_not(in_data.mask)]
        org_mask = np.logical_and(org_mask, in_data.mask)
        out_ds.write(org_data, window=Win)
    del in_ds,out_ds

if __name__ == '__main__':

    parser = ArgumentParser()
    parser.add_argument(
        '-p', '--path', help='the folder under which the files need mosaiced'
    )
    parser.add_argument(
        '-o', '--output_dir', help='the folder to save the output file'
    )
    parser.add_argument(
        '-f', '--prefix', default='', help='prefix of out fileName'
    )
    parser.add_argument(
        '-k', '--kw', default='', help='key word (means date) of spacenet'
    )
    args = parser.parse_args()

    #--------- Param list
    prefix = args.prefix
    kw = args.kw
    #-- 匹配串
    pattern = f'*{kw}*{kw}*.tif' 
    #-- 该文件夹下存放了待拼接的栅格
    path = args.path
    #-- 拼接结果输出文件
    output_path = os.path.join(args.output_dir, f"{prefix}_{kw}.tif")
    #-- 压缩结果输出文件
    compress_path = os.path.join(args.output_dir, f"{prefix}_compress_{kw}.tif")


    os.chdir(path)
    raster_list = sorted(glob(pattern)) #读取文件夹下所有tif数据
    print("-------------- {} tifs have been selected --------------".format(raster_list.__len__()))
    RasterMosaic(raster_list,outpath = output_path ) #拼接栅格
    compress(output_path,target_path = compress_path) #压缩栅格